#include "tbf.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <errno.h>
#include <signal.h>
#include <sys/time.h>
#include <pthread.h>

#define MAXTBF 1024

typedef struct tbf_st {
	int token; // 令牌
	int cps;   // 速率
	int burst; // 上限

	pthread_mutex_t mut;
	pthread_cond_t cond;
} tbf_t;

static pthread_mutex_t lib_mut = PTHREAD_MUTEX_INITIALIZER;
static tbf_t *libs[MAXTBF];
static int inited;

static int __get_free_pos(void)
{
	int i;

	pthread_mutex_lock(&lib_mut);
	for (i = 0; i < MAXTBF; i++) {
		if (NULL == libs[i]) {
			pthread_mutex_unlock(&lib_mut);
			return i;
		}
	}

	pthread_mutex_unlock(&lib_mut);

	return -1; // 满了
}

static void *thr_job(void *s)
{
	int td;
	tbf_t *cur;

	while (1) {
		pthread_mutex_lock(&lib_mut);
		for (td = 0; td < MAXTBF; td++) {
			cur = libs[td];
			if (NULL != cur) {
				pthread_mutex_lock(&cur->mut);
				cur->token += cur->cps;
				if (cur->token > cur->burst) {
					cur->token = cur->burst;
				}
				pthread_cond_broadcast(&cur->cond);
				pthread_mutex_unlock(&cur->mut);
			}
		}
		pthread_mutex_unlock(&lib_mut);

		sleep(1);
	}

	return (void *) 0;
}

static pthread_t tid;

static void module_load()
{
	int err;

	err = pthread_create(&tid, NULL, thr_job, NULL);

	if (err) {
		fprintf(stderr, "Error creating thread: %s\n", strerror(err));
		exit(1);
	}
}

int tbf_init(int cps, int burst)
{
	int td;
	tbf_t *new_tbf; // 创建令牌桶

	for (td = 0; td < MAXTBF; td++) {
		if (NULL == libs[td]) {
			break;
		}
	}
	// 开辟存储空间
	new_tbf = malloc(sizeof (tbf_t));
	// 成员初始化
	new_tbf->token = 0;
	new_tbf->cps = cps;
	new_tbf->burst = burst;
	pthread_mutex_init(&new_tbf->mut, NULL);
	pthread_cond_init(&new_tbf->cond, NULL);
	// 存放到令牌桶库中
	libs[td] = new_tbf;
	if (!inited) {
		module_load();
		inited = 1;
	}
	return td; // 返回初始化的桶下标
}

// 返回取到令牌的个数
int tbf_fetch_token(int td, int ntoken)
{
	if (!(td >= 0 && td < MAXTBF) || NULL == libs[td])
		return 0;

	pthread_mutex_lock(&libs[td]->mut);

	while (libs[td]->token < ntoken) { // 现在的token还不够取的，等一会
		pthread_cond_wait(&libs[td]->cond, &libs[td]->mut);
	}
	// 把token取出来，token数减少
	libs[td]->token -= ntoken;

	pthread_mutex_unlock(&libs[td]->mut);
	
	return ntoken; // 取出来这么多个
}

// 返回是否成功
int tbf_return_token(int td, int ntoken)
{
	if (NULL == libs[td])
		return -1;

	pthread_mutex_lock(&libs[td]->mut);

	// 把token放进去，token数增加
	libs[td]->token += ntoken;
	if (libs[td]->token > libs[td]->burst)
		libs[td]->token = libs[td]->burst;

	pthread_cond_broadcast(&libs[td]->cond);
	pthread_mutex_unlock(&libs[td]->mut);

	return 0;
}

void tbf_destroy(int td)
{
	pthread_mutex_lock(&lib_mut);

	if (NULL != libs[td]) {
		pthread_mutex_destroy(&libs[td]->mut);
		pthread_cond_destroy(&libs[td]->cond);

		free(libs[td]);

		libs[td] = NULL;
	}

	pthread_mutex_unlock(&lib_mut);
}

void tbf_destroy_all(void)
{
	int td;

	pthread_cancel(tid);

	for (td = 0; td < MAXTBF; td++) {
		if (NULL != libs[td]) {
			tbf_destroy(td);
		}
	}

	pthread_mutex_destroy(&lib_mut);
}

